import * as DOMPurify from 'dompurify';
import { JSDOM } from 'jsdom';
import * as vscode from 'vscode';

/**
 * View type that uniquely identifies the Mermaid chat output renderer.
 */
const viewType = 'vscode-samples.mermaid';

/**
 * Mime type used to identify Mermaid diagram data in chat output.
 */
const mime = 'application/vnd.chat-output-renderer.mermaid';

const maxFixAttempts = 3;

export function activate(context: vscode.ExtensionContext) {

	// Register our tools

	// The first tool takes mermaid markup as input
	context.subscriptions.push(
		vscode.lm.registerTool<{ markup: string }>('renderMermaidDiagram', {
			invoke: async (options, token) => {
				let sourceCode = options.input.markup;
				sourceCode = await runMermaidMarkupFixLoop(sourceCode, token);
				return writeMermaidToolOutput(sourceCode);
			},
		})
	);

	// The second tool generates mermaid markup based on a description
	context.subscriptions.push(
		vscode.lm.registerTool<{ description: string }>('createMermaidDiagram', {
			invoke: async (options, token) => {
				const description = options.input.description;

				let sourceCode = await generateMermaidDiagram(description, token);
				if (!sourceCode) {
					throw new Error('Failed to generate Mermaid diagram from description');
				}

				sourceCode = await runMermaidMarkupFixLoop(sourceCode, token);

				return writeMermaidToolOutput(sourceCode);
			},
		})
	);

	// Register the chat output renderer for Mermaid diagrams.
	// This will be invoked with the data generated by the tools.
	// It can also be invoked when rendering old Mermaid diagrams in the chat history.
	context.subscriptions.push(
		vscode.chat.registerChatOutputRenderer(viewType, {
			async renderChatOutput({ value }, webview, _ctx, _token) {
				const mermaidSource = new TextDecoder().decode(value);

				// Set the options for the webview
				const mermaidDist = vscode.Uri.joinPath(context.extensionUri, 'node_modules', 'mermaid', 'dist');
				webview.options = {
					enableScripts: true,
					localResourceRoots: [mermaidDist],
				};

				// Set the HTML content for the webview
				const nonce = getNonce();
				const mermaidEsmUri = vscode.Uri.joinPath(mermaidDist, 'mermaid.esm.mjs');

				webview.html = `
					<!DOCTYPE html>
					<html lang="en">

					<head>
						<meta charset="UTF-8">
						<meta name="viewport" content="width=device-width, initial-scale=1.0">
						<title>Mermaid Diagram</title>
						<meta http-equiv="Content-Security-Policy" content="default-src 'none'; script-src ${webview.cspSource} 'nonce-${nonce}'; style-src 'self' 'unsafe-inline';" />
					</head>

					<body>
						<pre class="mermaid">
							${escapeHtmlText(mermaidSource)}
						</pre>
						
						<script type="module" nonce="${nonce}">
							import mermaid from '${escapeForScriptBlock(webview.asWebviewUri(mermaidEsmUri).toString())}';
							mermaid.initialize({ startOnLoad: true });
						</script>
					</body>
					</html>`;
			},
		}));
}

/**
 * Lazily load mermaid
 */
const getMermaidInstance = (() => {
	const createMermaidInstance = async () => {
		// Patch the global window object for mermaid

		const { window } = new JSDOM("");
		// eslint-disable-next-line @typescript-eslint/no-explicit-any
		(global as any).window = window;
		// eslint-disable-next-line @typescript-eslint/no-explicit-any
		(global as any).DOMPurify = DOMPurify(window);
		return import('mermaid');
	};

	let cached: Promise<typeof import('mermaid')> | undefined;
	return async (): Promise<typeof import('mermaid').default> => {
		cached ??= createMermaidInstance();
		return (await cached).default;
	};
})();

/**
 * Tries to fix mermaid syntax errors in a set number of attempts.
 * 
 * @returns The best effort to fix the Mermaid markup.
 */
async function runMermaidMarkupFixLoop(sourceCode: string, token: vscode.CancellationToken): Promise<string> {
	let attempt = 0;
	while (attempt < maxFixAttempts) {
		const result = await validateMermaidMarkup(sourceCode);
		if (token.isCancellationRequested) {
			throw new Error('Operation cancelled');
		}

		if (result.type === 'success') {
			return sourceCode;
		}

		attempt++;

		sourceCode = await tryFixingUpMermaidMarkup(sourceCode, result.message, token);
		if (token.isCancellationRequested) {
			throw new Error('Operation cancelled');
		}
	}

	// Return whatever we have after max attempts
	return sourceCode;
}

/**
 * Validates the syntax of the provided Mermaid markup.
 */
async function validateMermaidMarkup(sourceCode: string): Promise<{ type: 'success' } | { type: 'error', message: string }> {
	try {
		const mermaid = await getMermaidInstance();
		await mermaid.parse(sourceCode);
		return { type: 'success' };
	} catch (error) {
		if (!(error instanceof Error)) {
			throw error;
		}

		return { type: 'error', message: error.message };
	}
}

/**
 * Uses a language model to try to fix Mermaid markup based on an error message.
 */
async function tryFixingUpMermaidMarkup(sourceCode: string, errorMessage: string, token: vscode.CancellationToken): Promise<string> {
	const model = await getPreferredLm();
	if (!model) {
		console.warn('No suitable model found for fixing Mermaid markup');
		return sourceCode;
	}

	if (token.isCancellationRequested) {
		throw new Error('Operation cancelled');
	}

	const completion = await model.sendRequest([
		vscode.LanguageModelChatMessage.Assistant(joinLines(
			`The user will provide you with the source code for the Mermaid diagram and an error message.`,
			`Your task is to fix the Mermaid source code based on the error message.`,
			`Please return the fixed Mermaid source code inside a \`mermaid\` fenced code block. Do not add any comments or explanation.`,
			`Make sure to return the entire source code.`
		)),
		vscode.LanguageModelChatMessage.User(joinLines(
			`Here is my Mermaid source code:`,
			``,
			`\`\`\`mermaid`,
			`${sourceCode}`,
			`\`\`\``,
			``,
			`And here is the mermaid error message:`,
			``,
			errorMessage,
		)),
	], {}, token);

	return await parseMermaidMarkupFromChatResponse(completion, token) ?? sourceCode;
}

async function parseMermaidMarkupFromChatResponse(chatResponse: vscode.LanguageModelChatResponse, token: vscode.CancellationToken): Promise<string | undefined> {
	const parts: string[] = [];
	for await (const line of chatResponse.text) {
		if (token.isCancellationRequested) {
			throw new Error('Operation cancelled');
		}

		parts.push(line);
	}

	const response = parts.join('');
	const lines = response.split('\n');
	if (!lines.at(0)?.startsWith('```') || !lines.at(-1)?.endsWith('```')) {
		console.warn('Invalid response format from model, expected fenced code block');
		return undefined;
	}

	return lines.slice(1, -1).join('\n').trim();
}

/**
 * Uses a language model to generate Mermaid markup based on a description of the diagram.
 */
async function generateMermaidDiagram(description: string, token: vscode.CancellationToken): Promise<string | undefined> {
	const model = await getPreferredLm();
	if (!model) {
		throw new Error('No suitable model found for generating Mermaid diagram');
	}

	const completion = await model.sendRequest([
		vscode.LanguageModelChatMessage.Assistant(joinLines(
			`The user will provide you with a description for a Mermaid diagram.`,
			`Your task is to generate the Mermaid source code based on the description.`,
			`Please return the Mermaid source code inside a \`mermaid\` fenced code block. Do not add any comments or explanation.`,
			`Make sure to return the entire source code.`,
		)),
		vscode.LanguageModelChatMessage.User(joinLines(
			`Please create a Mermaid diagram based on this description:`,
			`${description}`,
		)),
	], {}, token);

	return parseMermaidMarkupFromChatResponse(completion, token);

}

async function getPreferredLm(): Promise<vscode.LanguageModelChat | undefined> {
	return (await vscode.lm.selectChatModels({ family: 'gpt-4o-mini' })).at(0)
		?? (await vscode.lm.selectChatModels({ family: 'gpt-4o' })).at(0)
		?? (await vscode.lm.selectChatModels({})).at(0);
}

function writeMermaidToolOutput(sourceCode: string): vscode.LanguageModelToolResult {
	// Expose the source code as a tool result for the LM
	const result = new vscode.LanguageModelToolResult([
		new vscode.LanguageModelTextPart(sourceCode)
	]);

	// And store custom data in the tool result details to indicate that a custom renderer should be used for it.
	// In this case we just store the source code as binary data.

	// Add cast to use proposed API
	(result as vscode.ExtendedLanguageModelToolResult2).toolResultDetails2 = {
		mime,
		value: new TextEncoder().encode(sourceCode),
	};

	return result;
}

function joinLines(...lines: string[]): string {
	return lines.join('\n');
}

function escapeHtmlText(str: string): string {
	return str
		.replace(/&/g, '&amp;')
		.replace(/</g, '&lt;')
		.replace(/>/g, '&gt;')
		.replace(/"/g, '&quot;')
		.replace(/'/g, '&#39;');
}

function escapeForScriptBlock(str: string): string {
	return str
		.replace(/\\/g, '\\\\')
		.replace(/'/g, "\\'")
		.replace(/"/g, '\\"')
		.replace(/\r/g, '\\r')
		.replace(/\n/g, '\\n')
		.replace(/<\/script>/gi, '<\\/script>');
}

function getNonce() {
	let text = '';
	const possible = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789';
	for (let i = 0; i < 64; i++) {
		text += possible.charAt(Math.floor(Math.random() * possible.length));
	}
	return text;
}
